import numpy as np
from audio_processing.base_processor import AudioProcessor


class LimiterProcessor(AudioProcessor):
    """限幅器处理器 - 优化版本"""

    def __init__(self, parameters):
        super().__init__(parameters)
        self.threshold = parameters.get('threshold', 0.8)
        self.ratio = parameters.get('ratio', 4.0)
        self.attack = parameters.get('attack', 0.01)  # 秒
        self.release = parameters.get('release', 0.1)  # 秒

        # 状态变量 - 每个通道独立
        self.envelope = None
        self.gain = None

        # 预计算系数
        self.attack_coeff = 0.0
        self.release_coeff = 0.0
        self._update_coefficients()

    def set_sample_rate(self, sample_rate: int):
        super().set_sample_rate(sample_rate)
        self._update_coefficients()

    def set_input_channels(self, channels: int):
        """设置输入通道数并初始化状态变量"""
        super().set_input_channels(channels)
        self.envelope = np.zeros(channels)
        self.gain = np.ones(channels)

    def _update_coefficients(self):
        """更新攻击和释放系数"""
        if self.sample_rate > 0:
            self.attack_coeff = np.exp(-1.0 / (self.attack * self.sample_rate))
            self.release_coeff = np.exp(-1.0 / (self.release * self.sample_rate))

    def process(self, audio_data: np.ndarray) -> np.ndarray:
        """应用限幅效果 - 优化版本"""
        if not self.is_active:
            return audio_data

        # 快速检查：如果信号很小，直接返回
        if np.max(np.abs(audio_data)) < 1e-6:
            return audio_data

        # 确保数据形状正确
        audio_data = self._ensure_correct_shape(audio_data)

        output = np.zeros_like(audio_data)

        # 使用向量化操作替代逐样本循环
        input_level = np.abs(audio_data)

        # 更新包络 - 向量化
        new_envelope = np.where(
            input_level > self.envelope,
            self.attack_coeff * self.envelope + (1 - self.attack_coeff) * input_level,
            self.release_coeff * self.envelope + (1 - self.release_coeff) * input_level
        )

        # 防止包络为0，避免除零错误
        safe_envelope = np.where(new_envelope > 1e-10, new_envelope, 1e-10)

        # 计算增益减少 - 向量化
        excess = np.maximum(safe_envelope - self.threshold, 0)
        reduction = excess / self.ratio
        target_gain = np.where(
            safe_envelope > self.threshold,
            (self.threshold + reduction) / safe_envelope,
            1.0
        )

        # 平滑增益变化 - 向量化
        gain_mask = target_gain < self.gain
        new_gain = np.where(
            gain_mask,
            self.attack_coeff * self.gain + (1 - self.attack_coeff) * target_gain,
            self.release_coeff * self.gain + (1 - self.release_coeff) * target_gain
        )

        # 应用增益
        output = audio_data * new_gain

        # 更新状态
        self.envelope = new_envelope
        self.gain = new_gain

        return output

    def update_parameters(self, new_parameters):
        super().update_parameters(new_parameters)

        old_attack = self.attack
        old_release = self.release

        self.threshold = new_parameters.get('threshold', self.threshold)
        self.ratio = new_parameters.get('ratio', self.ratio)
        self.attack = new_parameters.get('attack', self.attack)
        self.release = new_parameters.get('release', self.release)

        # 如果时间常数改变，更新系数
        if old_attack != self.attack or old_release != self.release:
            self._update_coefficients()

    def reset(self):
        super().reset()
        if self.envelope is not None:
            self.envelope.fill(0.0)
        if self.gain is not None:
            self.gain.fill(1.0)